• Create ticket/issue

Follow the guide to enable codebase context checks.

Objective: To create a detailed and reliable record of critical system actions for security analysis
and compliance.

Status: Passed

Learn more about managing compliance generic rules or creating your own custom rules

Objective: To prevent the leakage of sensitive system information through error messages while
providing sufficient detail for internal debugging.

Status: Passed

Learn more about managing compliance generic rules or creating your own custom rules

Objective: To ensure logs are useful for debugging and auditing without exposing sensitive
information like PII, PHI, or cardholder data.

Status: Passed

Learn more about managing compliance generic rules or creating your own custom rules

Objective: Ensure all identifiers clearly express their purpose and intent, making code
self-documenting

Status:
Misleading parameter name: The constructor parameter n_samples is documented as "n events" but is used as
the embedding's vocabulary/num-embeddings size, which is misleading and harms
self-documentation.

def __init__(self, n_samples: int, emb_dim: int):
    super(VariableSortedHistoryPooling, self).__init__()
    # n samples are n events, where it's consecutive events belonging to a given user
    # The n samples can be segmented into B users.
    self.emb = torch.nn.Embedding(n_samples, emb_dim)

Learn more about managing compliance generic rules or creating your own custom rules

Objective: Ensure comprehensive error handling that provides meaningful context and graceful
degradation

Status:
Missing input validation: forward() performs no validation that offsets are sorted, within bounds, and consistent
with event_indices length (e.g., sum of user_lengths), which can cause runtime errors or
incorrect pooling on edge cases.

def forward(self, event_indices: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
    event_embs = self.emb(event_indices)
    # diffs of cumulative offsets gives user lengths (number of events in history per user)
    user_lengths = offsets[1:] - offsets[:-1]
    user_ids = torch.repeat_interleave(torch.arange(len(user_lengths),
                                                    device=offsets.device), 
                                        user_lengths)
    target = torch.zeros(len(user_lengths), event_embs.shape[1], device=event_embs.device)
    target = target.scatter_add(dim=0, index=user_ids.unsqueeze(1).expand_as(event_embs), src=event_embs)

Learn more about managing compliance generic rules or creating your own custom rules

Objective: Ensure all data inputs are validated, sanitized, and handled securely to prevent
vulnerabilities

Status:
Unvalidated tensor inputs: The module accepts event_indices and offsets without checking dtype/range constraints
(e.g., event_indices within embedding vocab, offsets monotonic/non-negative), and it is
unclear from the diff whether upstream guarantees make this safe.

def forward(self, event_indices: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
    event_embs = self.emb(event_indices)
    # diffs of cumulative offsets gives user lengths (number of events in history per user)
    user_lengths = offsets[1:] - offsets[:-1]
    user_ids = torch.repeat_interleave(torch.arange(len(user_lengths),
                                                    device=offsets.device), 
                                        user_lengths)
    target = torch.zeros(len(user_lengths), event_embs.shape[1], device=event_embs.device)
    target = target.scatter_add(dim=0, index=user_ids.unsqueeze(1).expand_as(event_embs), src=event_embs)
    return target / user_lengths.clamp(min=1).unsqueeze(1) 

Learn more about managing compliance generic rules or creating your own custom rules

Replace the custom VariableSortedHistoryPooling module with the standard
torch.nn.EmbeddingBag using mode='mean'. This change simplifies the
implementation to a single, optimized layer call.

Examples:

class VariableSortedHistoryPooling(torch.nn.Module):
    def __init__(self, n_samples: int, emb_dim: int):
        super(VariableSortedHistoryPooling, self).__init__()
        # n samples are n events, where it's consecutive events belonging to a given user
        # The n samples can be segmented into B users.
        self.emb = torch.nn.Embedding(n_samples, emb_dim)
    def forward(self, event_indices: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
        event_embs = self.emb(event_indices)
        # diffs of cumulative offsets gives user lengths (number of events in history per user)
        user_lengths = offsets[1:] - offsets[:-1]

 ... (clipped 6 lines)

Solution Walkthrough:

Before:

class VariableSortedHistoryPooling(torch.nn.Module):
    def __init__(self, n_samples: int, emb_dim: int):
        self.emb = torch.nn.Embedding(n_samples, emb_dim)
    
    def forward(self, event_indices: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
        event_embs = self.emb(event_indices)
        user_lengths = offsets[1:] - offsets[:-1]
        user_ids = torch.repeat_interleave(torch.arange(len(user_lengths)), user_lengths)
        target = torch.zeros(len(user_lengths), event_embs.shape[1], ...)
        target = target.scatter_add(0, user_ids.unsqueeze(1).expand_as(event_embs), event_embs)
        return target / user_lengths.clamp(min=1).unsqueeze(1)

After:

# The entire VariableSortedHistoryPooling class can be removed.
# It can be replaced by a single line in the model definition.

# In tests/test_pooling.py or any usage site:

# Before:
# model = VariableSortedHistoryPooling(vocab_size, emb_dim)

# After:
model = torch.nn.EmbeddingBag(vocab_size, emb_dim, mode='mean')

# The forward pass call remains the same and produces the same result:
out_vec = model(event_indices, offsets)

Add validation checks for the offsets tensor to ensure it is a 1D,
non-decreasing tensor with at least two elements.

LLMs/torch_examples/pooling.py [12]

+if offsets.dim() != 1 or offsets.size(0) < 2:
+    raise ValueError("offsets must be a 1D tensor with at least two elements")
+if not torch.all(offsets[1:] >= offsets[:-1]):
+    raise ValueError("offsets must be non-decreasing")
 user_lengths = offsets[1:] - offsets[:-1]

• Apply / Chat

Ensure the event_indices tensor is of torch.long dtype before passing it to the
embedding layer by casting it if necessary.

LLMs/torch_examples/pooling.py [10]

+if event_indices.dtype != torch.long:
+    event_indices = event_indices.long()
 event_embs = self.emb(event_indices)

• Apply / Chat

Rename the n_samples parameter in init to num_embeddings to accurately
reflect that it represents the vocabulary size, not the number of samples.

LLMs/torch_examples/pooling.py [3-8]

 class VariableSortedHistoryPooling(torch.nn.Module):
-    def __init__(self, n_samples: int, emb_dim: int):
+    def __init__(self, num_embeddings: int, emb_dim: int):
         super(VariableSortedHistoryPooling, self).__init__()
-        # n samples are n events, where it's consecutive events belonging to a given user
-        # The n samples can be segmented into B users.
-        self.emb = torch.nn.Embedding(n_samples, emb_dim)
+        # num_embeddings is the total number of unique events (i.e. vocabulary size).
+        self.emb = torch.nn.Embedding(num_embeddings, emb_dim)

• Apply / Chat